#include "benchmark_config.h"
#include "gpu_timer.h"
#include "random.h"
#include <bits/stdint-intn.h> // For `intN_t`.
#include <cstdlib>            // For `atoi`.
#include <iostream>
#include <stdio.h>
#include <thrust/device_vector.h>
#include <thrust/host_vector.h>
#include <thrust/reduce.h>

const int warmup_num = WARMUPNUM;
const int run_num = 1000;
const int element_value = 1;

template <typename T> void experiment_reduce(int size_scale) {

  // When the variable type is float, non-zero decimal parts may introduce
  // severe precision loss in calculations.
  // Using fixed values with zero decimal fraction avoids precision issues while
  // maintaining algorithm performance.
  thrust::host_vector<T> host_vector_input(size_scale, (T)element_value);
  thrust::device_vector<T> device_vector_input = host_vector_input;

  T cpu_result = 0;
  T gpu_result = 0;

  // warmup
  for (int w = 0; w < warmup_num; w++) {
    thrust::reduce(device_vector_input.begin(), device_vector_input.end());
  }
  // Ensure all warmup operations have completed before timing
  cudaDeviceSynchronize();

  GpuTimer timer;
  double total_time = 0;
  for (int i = 0; i < run_num; i++) {
    timer.start();
    // use thrust::reduce (default: addition)
    gpu_result =
        thrust::reduce(device_vector_input.begin(), device_vector_input.end());
    timer.stop_and_wait();
    total_time += timer.duration(1);
  }

  // Calculate performance and print
  double avg_time = total_time / run_num;
  const int64_t size_bytes = size_scale * sizeof(T);
  double performance = size_bytes * 1000. / 1024. / 1024. / 1024. / avg_time;
  printf("Average reduce algorithm performance: %f GB/s\n", performance);

  // verify
  cpu_result = size_scale * (T)element_value;
  bool correct = (gpu_result != cpu_result) ? false : true;
  printf("Thrust reduce %s\n", correct ? "PASS" : "FAIL");
}

int main(int argc, char *argv[]) {

  if (argc != 3) {
    printf("Usage: %s <data scale> <Data type: "
           "0(int32_t), 1(int64_t), 2(float), 3(double)>\n",
           argv[0]);
    return 1;
  }
  const int size_scale = atoi(argv[1]);
  const int type = atoi(argv[2]);

  switch (type) {
  case 0:
    printf("=========== Data type is int32_t ==========\n");
    experiment_reduce<int32_t>(size_scale);
    break;
  case 1:
    printf("=========== Data type is int64_t ==========\n");
    experiment_reduce<int64_t>(size_scale);
    break;
  case 2:
    printf("=========== Data type is float ==========\n");
    experiment_reduce<float>(size_scale);
    break;
  case 3:
    printf("=========== Data type is double ==========\n");
    experiment_reduce<double>(size_scale);
    break;
  default:
    break;
  }

  return 0;
}
